import xml.etree.ElementTree as ET
import os


class CreateXml:
    def __init__(self, folder, filename, path, img_size, label_list, database_src="Unknown"):
        self.folder = folder
        self.filename = filename
        self.path = path
        self.img_size = img_size
        self.database_src = database_src
        self.label_list = label_list
        self.dir = ""   #存储路径

        # 节点创建,根节点
        self.root = ET.Element("annotation")
        # 一级节点
        child_folder = ET.SubElement(self.root, "folder")
        child_filename = ET.SubElement(self.root, "filename")
        child_path = ET.SubElement(self.root, "path")
        child_source = ET.SubElement(self.root, "source")
        child_size = ET.SubElement(self.root, "size")
        child_segmented = ET.SubElement(self.root, "segmented")
        # 二级节点
        grandchild_database = ET.SubElement(child_source, "database")
        grandchild_width = ET.SubElement(child_size, "width")
        grandchild_height = ET.SubElement(child_size, "height")
        grandchild_depth = ET.SubElement(child_size, "depth")
        # 节点文本
        child_folder.text = self.folder
        child_filename.text = self.filename
        child_path.text = self.path
        grandchild_database.text = self.database_src
        grandchild_width.text = str(self.img_size[0])
        grandchild_height.text = str(self.img_size[1])
        grandchild_depth.text = str(self.img_size[2])
        child_segmented.text = str(0)

        self.object_generate()

    def object_generate(self):
        '''object标签下的信息'''
        def generate(parent, dictionary):
            for key, value in dictionary.items():
                if isinstance(value,dict):
                    child = ET.SubElement(parent,key)
                    generate(child,value)
                else:
                    child = ET.SubElement(parent,key)
                    child.text = str(value)
        for dic in self.label_list:
            child_object = ET.SubElement(self.root,'object')
            if isinstance(dic,dict):
                generate(child_object, dic)

    def save(self):
        '''保存xml'''
        self.indent(self.root)
        tree = ET.ElementTree(self.root)
        filename, extension = os.path.splitext(self.filename)
        if self.dir.strip():
            path_file = self.dir + '\\' + "{}.xml".format(filename)
            tree.write(path_file, encoding="utf-8")
        else:
            tree.write("{}.xml".format(filename), encoding="utf-8")

    def indent(self, elem, level=0):
        '''调整缩进'''
        i = "\n" + level * "    "
        if len(elem):
            if not elem.text or not elem.text.strip():#元素文本为空
                elem.text = i + "    "
            if not elem.tail or not elem.tail.strip():#tail指的是标签后面的文本信息，即不在标签内，而是标签外部，结束标签后面的位置
                elem.tail = i
            for elem in elem:
                self.indent(elem, level + 1)
            if not elem.tail or not elem.tail.strip():
                elem.tail = i
                '''在递归完之后，给下一级标签的最后一个标签尾部赋值，elem中储存着的是for循环中的最后一位，使得缩进与上一级一致
                因为是通过在尾部增加换行和缩进，因此会导致一级标签内的二级标签不断的在尾部添加二级缩进，直到最后一个二级标签的尾部是一级标签的结束标签，
                从而导致一级标签的结束标签与二级标签缩进一致。'''
        else:
            if level and (not elem.tail or not elem.tail.strip()):
                elem.tail = i

class ReadXml:
    '''根据输入的图片文件路径，寻找与文件同名的xml文件'''
    def __init__(self, file_path):
        #初始化变量
        self.picture_size = None
        self.item_label = []

        file_tempname = os.path.basename(file_path)
        file_name, extend = os.path.splitext(file_tempname)
        xml_name = "{}.xml".format(file_name)
        #current_path = os.path.dirname(os.getcwd())
        current_path = os.getcwd()
        xml_file = os.path.join(current_path,xml_name)

        if os.path.exists(xml_file):
            tree = ET.parse(xml_file)
            root = tree.getroot()
            size = root.find('size')
            ob = root.findall('object')

            width = size.find('width')
            height = size.find('height')
            depth = size.find('depth')

            self.picture_size = {'width': int(width.text), 'height': int(height.text), 'depth': int(depth.text)}

            for obj in ob:
                name = obj.find('name')
                bndbox = obj.find('bndbox')
                xmin = bndbox.find('xmin')
                ymin = bndbox.find('ymin')
                xmax = bndbox.find('xmax')
                ymax = bndbox.find('ymax')
                self.item_label.append({'name': name.text, 'xmin': xmin.text, 'ymin': ymin.text, 'xmax': xmax.text, 'ymax': ymax.text})

    def size(self):
        if self.picture_size:
            return self.picture_size

    def information(self):
        if self.item_label:
            return self.item_label


if __name__ == '__main__':
    '''test'''
    folder = "Picture"
    filename = "hohohoihiou"
    path = "C:\\Users\\HERO\\Pictures\\app图标"
    size = [200,390,3]
    dic_list = [{"name": "cat", "bndbox": {"xmin": "20", "ymin": "20", "xmax": "430", "ymax": "370"}},
                {"name": "gc", "bndbox": {"xmin": "32", "ymin": "45", "xmax": "490", "ymax": "370"}},
                {"name": "dog", "bndbox": {"xmin": "20", "ymin": "20", "xmax": "490", "ymax": "370"}}]
    read_xml = ReadXml("C:\\Users\\HERO\\Pictures\\20200411154104_14688.jpg")


